package com.lhf.aicloud.controller;

import com.baomidou.mybatisplus.core.conditions.query.LambdaQueryWrapper;
import com.baomidou.mybatisplus.core.conditions.update.LambdaUpdateWrapper;
import com.lhf.aicloud.entity.Answer;
import com.lhf.aicloud.entity.User;
import com.lhf.aicloud.entity.enums.AIModelEnum;
import com.lhf.aicloud.entity.enums.AITypeEnum;
import com.lhf.aicloud.mapper.UserMapper;
import com.lhf.aicloud.service.IAnswerService;
import com.lhf.aicloud.utils.*;
import jakarta.annotation.Resource;
import jakarta.servlet.http.HttpServletRequest;
import lombok.SneakyThrows;
import lombok.extern.slf4j.Slf4j;
import org.redisson.api.RLock;
import org.redisson.api.RedissonClient;
import org.springframework.ai.image.ImagePrompt;
import org.springframework.ai.image.ImageResponse;
import org.springframework.ai.openai.OpenAiChatClient;
import org.springframework.ai.openai.OpenAiImageClient;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.data.redis.core.RedisTemplate;
import org.springframework.util.StringUtils;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RestController;

import java.io.ByteArrayInputStream;
import java.io.InputStream;
import java.time.LocalDateTime;
import java.util.ArrayList;
import java.util.List;
import java.util.Objects;
import java.util.UUID;
import java.util.concurrent.TimeUnit;

/**
 * @author : lhf
 * @date : Created in 2024/07/27
 */
@RequestMapping("/openai")
@RestController
@Slf4j
public class OpenAIController {

    @Resource
    private UserMapper userMapper;

    @Resource
    private OpenAiChatClient openAiChatClient;

    @Resource
    private OpenAiImageClient openAiImageClient;

    @Autowired
    private IAnswerService answerService;

    @Resource
    private MinIoUtil minIoUtil;

    @Resource
    private RedissonClient redissonClient;

    @Resource
    private RedisTemplate<Object, Object> redisTemplate;

    /**
     * 调用 OpenAI 聊天接口
     *
     * @param question 问题
     */
    @RequestMapping("/chat")
    public ResponseEntity chat(String question, HttpServletRequest request) {
        if (!StringUtils.hasLength(question.trim())) {
            return ResponseEntity.error(400, "问题不能为空");
        }
        log.info("question: {}", question);

        String result = "";

        boolean addResult = false;

        // 获取当前用户ID
        Long uid = Objects.requireNonNull(SecurityUtils.getCurrentUser()).getUid();

        // 执行分布式锁
        String lockKey = AppVariable.getModelLockKey(uid, AIModelEnum.OPENAI.getValue(), AITypeEnum.CHAT.getValue());

        // 1. 获取分布式锁实例
        RLock rLock = redissonClient.getLock(lockKey);

        try {
            // 2. 获取分布式锁
            boolean isLock = rLock.tryLock(30, TimeUnit.SECONDS);
            if (!isLock) {
                return ResponseEntity.error("请求过于频繁，请稍后再试！");
            }

            // 查询当前用户当天可调用大模型的次数
            User user = userMapper.selectById(uid);
            if (user.getUsecount() <= 0) {
                return ResponseEntity.error("今日已无剩余调用次数，请明日再试！");
            }
            // 扣减当前用户当天调用大模型的次数
            user.setUsecount(user.getUsecount() - 1);
            userMapper.updateById(user);

            // 调用 OpenAI 聊天接口
            result = openAiChatClient.call(question);

            log.info("result: {}", result);

            // 将结果保存到 数据库
            Answer answer = new Answer();
            answer.setQuestion(question);
            answer.setAnswer(result);
            answer.setModel(AIModelEnum.OPENAI.getValue());
            answer.setCreatetime(LocalDateTime.now());
            answer.setUid(uid);
            answer.setType(AITypeEnum.CHAT.getValue());
            addResult = answerService.save(answer);

        } catch (Exception e) {
            throw new RuntimeException(e);
        } finally {
            // 3. 释放分布式锁
            rLock.unlock();
        }

        if (addResult) {
            return ResponseEntity.success(result);
        }
        return ResponseEntity.error("数据保存失败！");
    }

    /**
     * 调用 OpenAI 图片生成接口
     */
    @RequestMapping("/draw")
    public ResponseEntity draw(String question, HttpServletRequest request) {
        if (!StringUtils.hasLength(question.trim())) {
            return ResponseEntity.error(400, "问题不能为空");
        }
        log.info("prompt: {}", question);

        // 获取当前用户ID
        Long uid = Objects.requireNonNull(SecurityUtils.getCurrentUser()).getUid();

        String lockKey = AppVariable.getModelLockKey(uid, AIModelEnum.OPENAI.getValue(), AITypeEnum.DRAW.getValue());

        String imageUrl = null;
        boolean addResult = false;

        // 获取分布式锁实例
        RLock rLock = redissonClient.getLock(lockKey);

        try {
            // 获取分布式锁
            boolean isLock = rLock.tryLock(30, TimeUnit.SECONDS);
            if (!isLock) {
                return ResponseEntity.error("请求过于频繁，请稍后再试！");
            }

            // 调用 OpenAI 图片生成接口
            ImageResponse result = openAiImageClient.call(new ImagePrompt(question));
            String imageRemoteUrl = result.getResult().getOutput().getUrl();
            log.info("result: {}", imageRemoteUrl);

            // 将图片数据存储的MinIo中
            imageUrl = "";
            String imageName = "openai-draw-" + UUID.randomUUID().toString().replace("-", "") + ".png";
            try (InputStream inputStream = new ByteArrayInputStream(ImageUtils.getImageBytesFromUrl(imageRemoteUrl))) {
                // 上传到 minio
                imageUrl = minIoUtil.upload("openai/" + imageName, inputStream, "image/png");

            } catch (Exception e) {
                throw new RuntimeException(e);
            }
            log.info("imageUrl: {}", imageUrl);

            // 将结果保存到 数据库
            Answer answer = new Answer();
            answer.setQuestion(question);
            answer.setAnswer(imageUrl);
            answer.setModel(AIModelEnum.OPENAI.getValue());
            answer.setCreatetime(LocalDateTime.now());
            answer.setUid(uid);
            answer.setType(AITypeEnum.DRAW.getValue());
            addResult = answerService.save(answer);

        } catch (InterruptedException e) {
            throw new RuntimeException(e);
        } finally {
            rLock.unlock();
        }
        if (addResult) {
            return ResponseEntity.success(imageUrl);
        }
        return ResponseEntity.error("数据保存失败！");
    }

    /**
     * 获取OpenAI聊天历史信息
     */
    @RequestMapping("/chatHistory")
    public ResponseEntity chatHistory() {
        // 获取当前用户ID
        Long uid = Objects.requireNonNull(SecurityUtils.getCurrentUser()).getUid();
        String cacheKey = AppVariable.getListCacheKey(uid, AIModelEnum.OPENAI.getValue(), AITypeEnum.CHAT.getValue());
        ArrayList cacheList = (ArrayList) redisTemplate.opsForValue().get(cacheKey);
        List<Answer> list = null;
        if (cacheList == null || cacheList.isEmpty()) {
            // 构建查询条件
            LambdaQueryWrapper<Answer> answerLambdaQueryWrapper = new LambdaQueryWrapper<Answer>()
                    .eq(Answer::getUid, uid)
                    .eq(Answer::getType, AITypeEnum.CHAT.getValue())
                    .eq(Answer::getModel, AIModelEnum.OPENAI.getValue())
                    .orderByDesc(Answer::getCreatetime);
            // 查询
            list = answerService.list(answerLambdaQueryWrapper);
            redisTemplate.opsForValue().set(cacheKey, list, 1, TimeUnit.DAYS);
            return ResponseEntity.success(list);
        }
        return ResponseEntity.success(cacheList);
    }

    /**
     * 获取OpenAI绘画历史信息
     */
    @RequestMapping("/drawHistory")
    public ResponseEntity drawHistory() {
        // 获取当前用户ID
        Long uid = Objects.requireNonNull(SecurityUtils.getCurrentUser()).getUid();
        String cacheKey = AppVariable.getListCacheKey(uid, AIModelEnum.OPENAI.getValue(), AITypeEnum.DRAW.getValue());
        ArrayList cacheList = (ArrayList) redisTemplate.opsForValue().get(cacheKey);
        List<Answer> list = null;
        if (cacheList == null || cacheList.isEmpty()) {
            // 构建查询条件
            LambdaQueryWrapper<Answer> answerLambdaQueryWrapper = new LambdaQueryWrapper<Answer>()
                    .eq(Answer::getUid, uid)
                    .eq(Answer::getType, AITypeEnum.DRAW.getValue())
                    .eq(Answer::getModel, AIModelEnum.OPENAI.getValue())
                    .orderByDesc(Answer::getCreatetime);
            // 查询
            list = answerService.list(answerLambdaQueryWrapper);
            redisTemplate.opsForValue().set(cacheKey, list, 1, TimeUnit.DAYS);
            return ResponseEntity.success(list);
        }
        return ResponseEntity.success(cacheList);
    }
}
